package com.whz.governance.redislimiter.tool;

import com.whz.governance.redislimiter.command.LeakyBucketCommand;
import com.whz.governance.redislimiter.command.LeakyBucketResult;
import com.whz.governance.redislimiter.command.ScriptParams;
import lombok.experimental.UtilityClass;
import org.springframework.util.Assert;

import java.time.temporal.ChronoUnit;
import java.util.List;

/**
 * @author whz
 */
@UtilityClass
public class LeakyBucketUtils {
    private static final LeakyBucketCommand COMMAND = new LeakyBucketCommand();

    /**
     * 时间单位为秒的漏桶执行命令
     *
     * @param key             关键字
     * @param replenishRate   漏桶流速(单位秒)
     * @param burstCapacity   漏桶容量
     * @param requestedTokens 每个请求的token个数,一般是1
     * @return 执行结果
     */
    public LeakyBucketResult execute(String key, int replenishRate, int burstCapacity, int requestedTokens) {
        return doExecute(key, replenishRate, burstCapacity, requestedTokens, ChronoUnit.SECONDS);
    }

    private LeakyBucketResult doExecute(String key, int replenishRate, int burstCapacity, int requestedTokens,
                                        ChronoUnit unit) {
        Assert.hasText(key, "leak bucket: key not null");
        Assert.isTrue(replenishRate > 0, "leak bucket: replenishRate > 0");
        Assert.isTrue(burstCapacity > replenishRate, "leak bucket:burstCapacity > replenishRate");
        Assert.isTrue(requestedTokens > 0, "leak bucket:requestedTokens > 0");
        ScriptParams params = ScriptParams.create()
                                          .addKey(key + "_token")
                                          .addKey(key + "_timestamp")
                                          .addArg(replenishRate + "")
                                          .addArg(burstCapacity + "")
                                          .addArg(InstantUtils.currentUnitDuration(unit) + "")
                                          .addArg(requestedTokens + "")
                                          // 设置redis超时时间
                                          .addArg(InstantUtils.getTtl(replenishRate, burstCapacity, unit) * 2 + "");
        List<Long> execute = COMMAND.execute(params);
        Assert.isTrue(execute != null && execute.size() == 2, "LeakyBucketCommand MUST return 2 long num.");
        LeakyBucketResult result = new LeakyBucketResult();
        result.setPermit(execute.get(0) == 1);
        result.setRequestTokens(execute.get(1));
        return result;
    }


    /**
     * 漏桶执行命令
     *
     * @param key             关键字
     * @param replenishRate   漏桶流速(单位秒)
     * @param burstCapacity   漏桶容量
     * @param requestedTokens 每个请求的token个数,一般是1
     * @param unit            指定的时间单位
     * @return 执行结果
     */
    public LeakyBucketResult execute(String key, int replenishRate, int burstCapacity, int requestedTokens,
                                     ChronoUnit unit) {
        return doExecute(key, replenishRate, burstCapacity, requestedTokens, unit);
    }
}